"""
Author: Morphlng
Date: 2024-04-01 11:00:00
LastEditTime: 2024-12-16 23:13:30
LastEditors: Morphlng
Description: RLAgent is an interactive entity that interacts with the environment
FilePath: /DrivingGym/src/driving_gym/environment/agent/rl_agent.py
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.action import ActionHandler
from driving_gym.environment.agent.done import DoneHandler
from driving_gym.environment.agent.obs import ObsHandler
from driving_gym.environment.agent.reward import RewardHandler
from driving_gym.environment.scenario.base_scenario import BaseScenario

if TYPE_CHECKING:
    from driving_gym.simulation.adapter_interface import AdapterInterface


class RLAgent:
    def __init__(self, config: dict, adapter: "AdapterInterface"):
        self.config = config
        self.actor_id = config["actor_id"]
        self._validate_config()

        self.obs_handler = ObsHandler(config["obs"], adapter)
        self.action_handler = ActionHandler(config["action"], adapter)
        self.reward_handler = RewardHandler(config["reward"], adapter)
        self.done_handler = DoneHandler(config["done"], adapter)

        self._action_space = None
        self._observation_space = None
        self._scenario = None

    @property
    def action_space(self):
        if self._action_space is None:
            self._action_space = self.action_handler.get_action_space()

        return self._action_space

    @property
    def observation_space(self):
        if self._observation_space is None:
            obs_space = self.obs_handler.get_observation_space()
            mask_space = self.action_handler.get_mask_space()
            if mask_space is not None:
                spaces = {**obs_space.spaces, "action_mask": mask_space}
                self._observation_space = type(obs_space)(spaces)
            else:
                self._observation_space = obs_space

        return self._observation_space

    def reset(self, scenario: BaseScenario):
        self._scenario = scenario
        self.obs_handler.reset()
        self.action_handler.reset()
        self.reward_handler.reset()
        self.done_handler.reset()

    def act(self, action, time_step):
        """Perform one step of action"""
        abstract_action = self.action_handler.convert_single_action(action, time_step)
        if abstract_action.done():
            return None
        else:
            return abstract_action.run_step()

    def get_obs(self, snapshots: list[Snapshot]):
        obs = self.obs_handler.get_obs(snapshots, self._scenario)
        action_mask = self.action_handler.get_action_mask()
        if action_mask is not None:
            obs["action_mask"] = action_mask
        return obs

    def get_reward(self, snapshot: Snapshot):
        reward, reward_info = self.reward_handler.get_reward(snapshot, self._scenario)
        return reward, reward_info

    def get_done(self, snapshot: Snapshot):
        done, done_info = self.done_handler.get_done(snapshot, self._scenario)
        return done, done_info

    def _validate_config(self):
        for key in ["obs", "action", "reward", "done"]:
            if key not in self.config:
                raise ValueError(f"Missing `{key}` in agent {self.actor_id}'s config")

        if "items" not in self.config["obs"]:
            raise ValueError(f"Missing `items` in agent {self.actor_id}'s obs config")

        for obs_type, obs_config in self.config["obs"]["items"].items():
            obs_config["actor_id"] = self.config["actor_id"]

        for key in ["obs", "action", "reward", "done"]:
            self.config[key]["actor_id"] = self.config["actor_id"]
